/* libwatopt++ - Optimization Toolkit
 *
 * Copyright (c) 2007 Christopher Alexander Watford
 * <christopher.watford@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 *
 * $Id: art2_network.cpp 2 2007-05-06 18:01:15Z christopher.watford $
 */

#include <vector>
#include <cmath>
#include <ctime>
#include <limits>

#include <art2_network.hpp>
#include "utils.hpp"

namespace watopt
{
   unsigned long start_time = (unsigned long)std::time(NULL);
   void visualize(const char*, double*, unsigned int, unsigned long, unsigned int);

   static size_t
      _zji_index(size_t jj, size_t ii, size_t categories)
   {
      return (categories*jj) + ii;
   }

   static size_t 
      _zij_index(size_t ii, size_t jj, size_t inputs)
   {
      return (inputs*ii) + jj;
   }

   ART2Network::ART2Network(
	   unsigned int inputs,
	   unsigned int categories,
	   double epsilon,
	   double a,
	   double b,
	   double c,
	   double d,
	   double theta,
      double rho)
	   : _f1_nodes(inputs),
	     _f2_nodes(categories),
	     _eps(epsilon),
	     _aa(a),
	     _bb(b),
	     _cc(c),
	     _dd(d),
	     _theta(theta),
        _rho(rho),
	     _leader(0),
        _ww_norm(0.0),
        _vv_norm(0.0),
        _pp_norm(0.0),
        _uu_norm(0.0),
        _rr_norm(0.0),
	     _ww(new double[inputs]),
	     _xx(new double[inputs]),
	     _vv(new double[inputs]),
	     _uu(new double[inputs]),
	     _pp(new double[inputs]),
	     _qq(new double[inputs]),
	     _rr(new double[inputs]),
        _yy(new double[categories]),
        _Tj(new double[categories]),
	     _reset(new bool[categories]),
        _zji(new double[categories*inputs]),
        _zij(new double[inputs*categories])
   {
	   for(unsigned int jj = 0; jj < categories; jj++)
	   {
         _yy[jj] = 0.0;
         for(unsigned int ii = 0; ii < inputs; ii++)
         {
            _zji[_zji_index(jj,ii,categories)] = 0.0;
         }
	   }

      double _zij_init = (1.0 / (((1.0 - _dd)*std::sqrt((double)inputs))));
	   for(unsigned int ii = 0; ii < inputs; ii++)
	   {
		   for(unsigned int jj = 0; jj < categories; jj++)
		   {
            _zij[_zij_index(ii,jj,inputs)] = _zij_init;
		   }
	   }

      std::memset(_ww, 0, sizeof(double)*inputs);
      std::memset(_xx, 0, sizeof(double)*inputs);
      std::memset(_vv, 0, sizeof(double)*inputs);
      std::memset(_uu, 0, sizeof(double)*inputs);
      std::memset(_pp, 0, sizeof(double)*inputs);
      std::memset(_qq, 0, sizeof(double)*inputs);
      std::memset(_rr, 0, sizeof(double)*inputs);
      std::memset(_Tj, 0, sizeof(double)*categories);// p_i initializes to 0.0, so the summed values will be 0 for p_i * z_ij
      std::memset(_reset, 0, sizeof(bool)*categories);
   }

   ART2Network::~ART2Network(void)
   {
 	     delete _ww;
        delete _xx;
        delete _vv;
        delete _uu;
        delete _pp;
        delete _qq;
        delete _rr;
        delete _Tj;
        delete _reset;
        delete _yy;
        delete _zij;
        delete _zji;
   }

   static double
   normalize(double* vec, unsigned int size, double cc = 1.0)
   {
      double sqsum = 0.0;
      for(unsigned int ii = 0; ii < size; ii++)
         sqsum += vec[ii]*vec[ii];

      return cc*std::sqrt(sqsum);
   }

   static double
   normalize(std::vector<double> vec)
   {
      double sqsum = 0.0;
      std::vector<double>::size_type size = vec.size();
      for(std::vector<double>::size_type ii = 0; ii < size; ii++)
         sqsum += vec[ii]*vec[ii];

      return std::sqrt(sqsum);
   }

   unsigned int
   ART2Network::ApplyInput(double* Inputs, unsigned int patternId)
   {
      visualize("I", Inputs, _f1_nodes, start_time, patternId);

      do
      {
         _f2_active = false;

	     //print_network(std::cout);

         // Run F1
         equilibrate(Inputs);

	     //print_network(std::cout);

         // Do F1->F2
         // Find the winner
         find_max();

	     _f2_active = true;

	     //print_network(std::cout);

         // Do F2->F1
         equilibrate(Inputs);

	     //print_network(std::cout);

         // Update r
         do_r();

         // Check for a reset and continue if so
      } while(check_reset());

      // Updated z values
      apply_dz();

      print_network(std::cout, patternId);

      // Clear F1
      std::memset(_ww, 0, sizeof(double)*_f1_nodes);
      std::memset(_xx, 0, sizeof(double)*_f1_nodes);
      std::memset(_vv, 0, sizeof(double)*_f1_nodes);
      std::memset(_uu, 0, sizeof(double)*_f1_nodes);
      std::memset(_pp, 0, sizeof(double)*_f1_nodes);
      std::memset(_qq, 0, sizeof(double)*_f1_nodes);

      // Clear reset
      std::memset(_reset, 0, sizeof(bool)*_f2_nodes);

      return _leader;
   }

   void
   ART2Network::equilibrate(double *_II)
   {
      bool first = true;
      double lastnp = _pp_norm;
      double dnp = 0.0; // delta ||p||

      while(first || (dnp >= _eps))
      {
      // 1. compute w_i
         do_w(_II);

      // 2. compute x_i
         do_x();

      // 3. compute v_i
         do_v();

      // 4. compute u_i
         do_u();

      // 5. compute p_i
         do_p();

         // our test for equilibrium... may need work
         dnp = std::abs(lastnp - _pp_norm);
         lastnp = _pp_norm;

      // 6. compute q_i
         do_q();

         if(first)
         {  // hack to make sure dnp is ignored first time around
            first = false;
         }
      }
   }

   void
   ART2Network::find_max(void)
   {
      // 1. do T_j, picks a leader
      do_T();
   }

   /*
    * Returns 'true' if there was a reset
    *
    */
   bool
   ART2Network::check_reset(void)
   {
      if((_rho / (_rr_norm+_eps)) > 1.0)
      {
         _reset[_leader] = true;
         return true; // reset, do not continue
      }

      return false; // we've been plenty vigilante
   }

   void
   ART2Network::apply_dz(void)
   {
      double tsse;
      do
      {
         tsse = 0.0;

         // 1. do bottom-up/top-down dz
         for(unsigned int ii = 0; ii < _f1_nodes; ii++)
         {
            _zij[_zij_index(ii,_leader,_f1_nodes)] += _dd * (_pp[ii] - _zij[_zij_index(ii,_leader,_f1_nodes)]);
            _zji[_zji_index(_leader,ii,_f2_nodes)] += _dd * (_pp[ii] - _zji[_zji_index(_leader,ii,_f2_nodes)]);

            tsse += (_zij[_zij_index(ii,_leader,_f1_nodes)]-_zji[_zji_index(_leader,ii,_f2_nodes)])
               *(_zij[_zij_index(ii,_leader,_f1_nodes)]-_zji[_zji_index(_leader,ii,_f2_nodes)]);
         }
      } while(tsse > _eps);
   }

   void
   ART2Network::do_w(double *_II)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _ww[ii] = _II[ii] + (_aa * _uu[ii]);
      }

      _ww_norm = normalize(_ww, _f1_nodes);
   }

   void
   ART2Network::do_x(void)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _xx[ii] = _ww[ii]/(_ww_norm+_eps);
      }
   }

   double
   ART2Network::f(double x)
   {
      return (x >= _theta ? x : 0.0);
   }

   void
   ART2Network::do_v(void)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _vv[ii] = f(_xx[ii]) + (_bb * f(_qq[ii]));
      }

      _vv_norm = normalize(_vv, _f1_nodes);
   }

   void
   ART2Network::do_u(void)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _uu[ii] = _vv[ii]/(_vv_norm+_eps);
      }

      _uu_norm = normalize(_uu, _f1_nodes);
   }

   void
   ART2Network::do_T(void)
   {
      double maxval = std::numeric_limits<double>::min();
      for(unsigned int jj = 0; jj < _f2_nodes; jj++)
      {
         _Tj[jj] = 0.0;
         for(unsigned int ii = 0; ii < _f1_nodes; ii++)
         {
            _Tj[jj] += _pp[ii]*_zij[_zij_index(ii,jj,_f1_nodes)]*(_reset[jj] ? 0.0 : 1.0);
         }

         if(_Tj[jj] > maxval)
         {
            _leader = jj;
            maxval = _Tj[jj];
         }
      }
   }

   double
   ART2Network::find_zJi(unsigned int ii)
   {
      if(!_f2_active || _reset[_leader])
         return 0.0;

      return _dd*_zji[_zji_index(_leader,ii,_f2_nodes)];
   }

   void
   ART2Network::do_p(void)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _pp[ii] = _uu[ii] + find_zJi(ii);
      }

      _pp_norm = normalize(_pp, _f1_nodes);
   }

   void
   ART2Network::do_q(void)
   {
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _qq[ii] = _pp[ii]/(_pp_norm+_eps);
      }   
   }

   void
   ART2Network::do_r(void)
   {
      double norm_u_cp = _uu_norm + _cc*_pp_norm + _eps;
      for(unsigned int ii = 0; ii < _f1_nodes; ii++)
      {
         _rr[ii] = (_uu[ii] + _cc*_pp[ii]) / norm_u_cp;
      }

      _rr_norm = normalize(_rr, _f1_nodes);
   }

   void
   visualize(const char* name, double* nodes, unsigned int node_count, unsigned long tt, unsigned int cnt)
   {
   #ifdef ART2__SELF_TEST
   #define IO_COPY_BLKSZ (4096)
      static char buffer[256];
      static char copyBlock[IO_COPY_BLKSZ];

      buffer[_snprintf(buffer, 255, "%ld-%d-%s.bmp", tt, cnt, name)] = '\0';

      FILE *tmpl = fopen("template.bmp", "rb");

      node_count = (unsigned int)std::sqrt((double)node_count);

      short header_offset = 0, bpp = 0; 

      fseek(tmpl, 10, SEEK_SET);
      fread(&header_offset, sizeof(short), 1, tmpl);

      fseek(tmpl, 28, SEEK_SET);
      fread(&bpp, sizeof(short), 1, tmpl);
      bpp /= 8;

      FILE *img = fopen(buffer, "wb");

      fseek(tmpl, 0, SEEK_SET);
      fseek(img, 0, SEEK_SET);

      // copy the header
	   size_t blocks = (size_t)header_offset / IO_COPY_BLKSZ;
	   size_t tail = (size_t)header_offset % IO_COPY_BLKSZ;
	   while(blocks)
	   {
		   if(fread(copyBlock, 1, IO_COPY_BLKSZ, tmpl) != IO_COPY_BLKSZ)
         {
            fclose(img);
			   return;
         }

		   if(fwrite(copyBlock, 1, IO_COPY_BLKSZ, img) != IO_COPY_BLKSZ)
         {
            fclose(tmpl);
			   return;
         }

		   blocks--;
	   }

	   if(tail)
	   {
		   if(fread(copyBlock, 1, tail, tmpl) != tail)
         {
            fclose(img);
			   return;
         }

		   if(fwrite(copyBlock, 1, tail, img) != tail)
         {
            fclose(tmpl);
			   return;
         }
	   }

      // write the data
      for(unsigned int ii = 0; ii < node_count*20; ii++)
      {
         for(unsigned int jj = 0; jj < node_count*20; jj++)
         {
            unsigned char byte = (unsigned char)(nodes[ii/20*node_count+jj/20] * 255);
            for(unsigned int bb = 0; bb < bpp; bb++)
            {
               fwrite(&byte, sizeof(unsigned char), 1, img);
            }
         }

         // rows are padded to 4
         if((node_count*20*bpp) % 4)
         {
            fseek(img, 4 - ((node_count*20*bpp)%4), SEEK_CUR);
         }
      }

      fclose(tmpl);
      fclose(img);
   #endif /* defined(ART2__SELF_TEST) */
   }

   #define PRINT_F1_NODE(node, cvar) \
	   os << node << std::endl; \
	   for(unsigned int ii = 0; ii < _f1_nodes; ii++) \
		   os << cvar[ii] << "\t"; \
	   os << std::endl

   void
   ART2Network::print_network(std::ostream &os, unsigned int pattern)
   {
   #ifdef ART2__SELF_TEST
      visualize("w", _ww, _f1_nodes, start_time, pattern);
      visualize("x", _xx, _f1_nodes, start_time, pattern);
      visualize("v", _vv, _f1_nodes, start_time, pattern);
      visualize("u", _uu, _f1_nodes, start_time, pattern);
      visualize("p", _pp, _f1_nodes, start_time, pattern);
      visualize("q", _qq, _f1_nodes, start_time, pattern);

	   return;

	   os << "==========================" << std::endl;
	   os << "F1 Layer (" << _f1_nodes << ")" << std::endl;
	   os << "==========================" << std::endl;
	   PRINT_F1_NODE("w", _ww);
	   PRINT_F1_NODE("x", _xx);
	   PRINT_F1_NODE("v", _vv);
	   PRINT_F1_NODE("u", _uu);
	   PRINT_F1_NODE("p", _pp);
	   PRINT_F1_NODE("q", _qq);

	   os << "==========================" << std::endl;
	   os << "F2 Layer (" << _f2_nodes << ")" << std::endl;
	   os << "==========================" << std::endl;
	   if(_f2_active)
	   {
		   os << "Top Down" << std::endl;
		   os << "--------------------------" << std::endl;
		   for(unsigned int jj = 0; jj < _f2_nodes; jj++)
		   {
			   for(unsigned int ii = 0; ii < _f1_nodes; ii++)
			   {
				   os << _zji[jj][ii] << "\t";
			   }
			   os << std::endl;
		   }
		   os << "--------------------------" << std::endl;
		   os << "Bottom Up" << std::endl;
		   os << "--------------------------" << std::endl;
		   for(unsigned int ii = 0; ii < _f1_nodes; ii++)
		   {
			   for(unsigned int jj = 0; jj < _f2_nodes; jj++)
			   {
				   os << _zij[ii][jj] << "\t";
			   }
			   os << std::endl;
		   }
		   os << "--------------------------" << std::endl;
		   os << "Reset" << std::endl;
		   os << "--------------------------" << std::endl;
		   for(unsigned int jj = 0; jj < _f2_nodes; jj++)
		   {
			   os << _reset[jj] << "\t";
		   }
		   os << std::endl;
		   os << "--------------------------" << std::endl;
		   os << "Leading Category" << std::endl;
		   os << "--------------------------" << std::endl;
		   os << _leader << std::endl;
	   } else {
		   os << "INACTIVE" << std::endl;
	   }

	   int x;
	   std::cin >> x;

	   os << std::endl;
	   return;
   #endif /* defined(ART2__SELF_TEST) */
   }

}
